//! A oneshot channel is a synchronization primitive used to send a single
//! message from one task to another. When the sender has sent a message and it
//! is ready to be received, the oneshot notifies the receiver. The receiver is
//! a Rust async `Future` that can be `await`ed. Our implementation does not
//! depend on std or alloc.
//!
//! See https://tweedegolf.nl/blog/50/build-your-own-async-primitive
//! for a full description of the internals.

use futures::task::AtomicWaker;
use std::cell::UnsafeCell;
use std::future::Future;
use std::mem::MaybeUninit;
use std::pin::Pin;
use std::sync::atomic::{AtomicBool, Ordering};
use std::task::{Context, Poll};

pub struct Oneshot<T> {
    message: UnsafeCell<MaybeUninit<T>>,
    waker: AtomicWaker,
    has_message: AtomicBool,
}
unsafe impl<T> Sync for Oneshot<T> {}

impl<T> Oneshot<T> {
    pub const fn new() -> Self {
        Self {
            message: UnsafeCell::new(MaybeUninit::uninit()),
            waker: AtomicWaker::new(),
            has_message: AtomicBool::new(false),
        }
    }

    pub unsafe fn put(&self, message: T) {
        self.message.get().write(MaybeUninit::new(message));
        self.has_message.store(true, Ordering::Release);
        self.waker.wake();
    }

    pub unsafe fn take(&self) -> Option<T> {
        if self.has_message.load(Ordering::Acquire) {
            let message = self.message.get().read().assume_init();
            self.has_message.store(false, Ordering::Release);
            Some(message)
        } else {
            None
        }
    }

    pub fn split<'a>(&'a mut self) -> (Sender<'a, T>, Receiver<'a, T>) {
        unsafe { self.take() };
        (Sender(self), Receiver(self))
    }

    pub unsafe fn recv<'a>(&'a self) -> Receiver<'a, T> {
        Receiver(self)
    }

    pub unsafe fn send<'a>(&'a self) -> Sender<'a, T> {
        Sender(self)
    }

    pub fn is_empty(&self) -> bool {
        !self.has_message.load(Ordering::AcqRel)
    }
}

impl<T> Drop for Oneshot<T> {
    fn drop(&mut self) {
        unsafe { self.take() };
    }
}

pub struct Sender<'a, T>(&'a Oneshot<T>);

impl<'a, T> Sender<'a, T> {
    pub fn send(self, message: T) {
        unsafe { self.0.put(message) };
    }
}

pub struct Receiver<'a, T>(&'a Oneshot<T>);

impl<'a, T> Future for Receiver<'a, T> {
    type Output = T;

    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        self.0.waker.register(cx.waker());

        if let Some(message) = unsafe { self.0.take() } {
            Poll::Ready(message)
        } else {
            Poll::Pending
        }
    }
}

impl<'a, T> Drop for Receiver<'a, T> {
    fn drop(&mut self) {
        unsafe { self.0.take() };
    }
}
